"""
    2048 Environment,在4bits版本基础上优化的5bits版本。
    棋盘状态用16个uint16类型数表示，其中每个数表示一个tile，
    与4bits版本不同的是5bits版本中每个tile的访问为0到16，即17个数，
    也就是GUI版本中的0， 2**1=2, ..., 2**15=32768, 2**16=65536
    board=np.array([[0, 0, 0, 0],
                    [0, 1, 0, 0],
                    [0, 0, 0, 0], 
                    [0, 0, 1, 0]], dtype=np.uint16) 
"""

import numpy as np


class Env_2048():
    def __init__(self) -> None:
        self.init_trans_table()

    def init_trans_table(self):
        """ init environment transition table"""
        self.move_left_table = np.zeros((17, 17, 17, 17, 4), dtype=np.uint16)
        self.move_right_table = np.zeros((17, 17, 17, 17, 4), dtype=np.uint16)
        self.score_table = np.zeros((17, 17, 17, 17), dtype=np.uint32)

        for a in range(17):
            for b in range(17):
                for c in range(17):
                    for d in range(17):
                        line = self.move_row([a, b, c, d])  # move left
                        self.move_left_table[a, b, c, d][::] = line
                        self.score_table[a, b, c, d] = self.score_line([a, b, c, d]) 

                        self.move_right_table[d, c, b, a][::] = line[::-1]  # move right

    def move_row(self, line):
        i = 0
        while(i<3):
            for j in range(i+1, 4):
                if line[j] != 0:
                    break
                j += 1
            if(j==4):
                break

            if(line[i]==0):
                line[i] = line[j]
                line[j] = 0
                if(j==3):
                    break
                else:
                    i -= 1
            elif(line[i]==line[j]):
                line[i] += 1
                line[j] = 0
            
            i += 1
        return line

    def score_line(self, line):
        score = 0

        for tile in line:
            if tile>=2:
                score += (tile-1)*(2**tile)
        return score

    def move_left(self, board:np.array):
        a = board[:, 0]
        b = board[:, 1]
        c = board[:, 2]
        d = board[:, 3]

        new_board = self.move_left_table[a, b, c, d]
        return new_board

    def move_right(self, board:np.array):
        a = board[:, 0]
        b = board[:, 1]
        c = board[:, 2]
        d = board[:, 3]

        new_board = self.move_right_table[a, b, c, d]
        return new_board

    def move_up(self, board:np.array):
        t_board = self.move_left(board.T)
        return t_board.T

    def move_down(self, board:np.array):
        t_board = self.move_right(board.T)
        return t_board.T

    def execute_move(self, move, board):
        if(move==0):
            return self.move_up(board)
        if(move==1):
            return self.move_down(board)
        if(move==2):
            return self.move_left(board)
        if(move==3):
            return self.move_right(board)

    def init_board(self):
        board = np.zeros(16, dtype=np.uint16)
        index_a, index_b = np.random.choice(16, 2, replace=False)
        tile_a, tile_b = np.random.choice([1, 2], size=2, p=[0.9, 0.1], replace=True)
        board[index_a], board[index_b] = tile_a, tile_b
        penalty = ((tile_a>>1) + (tile_b>>1))*4
        return board.reshape((4, 4)), penalty

    def insert_tile(self, board):
        """
            board=np.array([[0, 0, 0, 0],
                            [0, 1, 0, 0],
                            [0, 0, 0, 0], 
                            [0, 0, 1, 0]], dtype=np.uint16) 
        """
        _board = np.copy(board)
        index_a, index_b = np.where(_board==0)
        choice_index = np.random.choice(len(index_a), 1, replace=False)

        tile = np.random.choice([1, 2], size=1, p=[0.9, 0.1], replace=True)
        _board[index_a[choice_index], index_b[choice_index]] = tile
        penalty = (tile>>1)*4
        return _board, int(penalty)

    def get_score(self, board, penalty):
        a = board[:, 0]
        b = board[:, 1]
        c = board[:, 2]
        d = board[:, 3]

        score = np.sum( self.score_table[a, b, c, d] )
        return score-penalty

    def count_distinct_tiles(self, board):
        """
            计算棋盘不同数值的个数
            假设棋盘没有死局，必定存在0数值，这里不计算0数值
        """
        return np.unique(board).size - 1

    def count_empty(self, board):
        return 16 - np.count_nonzero(board)

    def available_move(self, board:np.array):
        up_board    = self.move_up(board) 
        down_board  = self.move_down(board)
        left_board  = self.move_left(board)
        right_board = self.move_right(board)

        not_available = np.array([
                    np.all(board==up_board),
                    np.all(board==down_board),
                    np.all(board==left_board),
                    np.all(board==right_board),
                    ], dtype=np.bool_)
        available = np.logical_not(not_available) 
        return available, [up_board, down_board, left_board, right_board]


def main():
    env = Env_2048()
    move_num = 0
    penalty = 0
    board, p = env.init_board()
    penalty += p

    while (True):
        score = env.get_score(board, penalty)
        print(board)
        print("move num: ", move_num, "  score: ", score)

        direct = input("please input:")
        if(direct == "a"):
            _board = env.move_left(board)
        if(direct == "d"):
            _board = env.move_right(board)
        if(direct == "w"):
            _board = env.move_up(board)
        if(direct == "s"):
            _board = env.move_down(board)
        if(np.all(_board == board)):
            print("无效移动，棋盘状态未变化!!!")
        else:
            move_num += 1
            board, p = env.insert_tile(_board)
            penalty += p


if __name__ == '__main__':
    main()
